# -*- coding: utf-8 -*-
# @Time     : 2020/11/10 2:31 下午
# @Author   : huangxiong
# @FileName : stock_info.py
# @Comment  :
# @Software : PyCharm

"""
脚本说明: 获取股票基本信息、财务数据
"""

from quant_researcher.quant.datasource_fetch.index_api.index_constant import SW_OLD_NEW_MAPPING, \
    SW_INDUSTRY_CODE_SNAME
from quant_researcher.quant.project_tool import time_tool
from quant_researcher.quant.project_tool.celebrity import get_lst_condition
from quant_researcher.quant.project_tool.db_operator import db_conn, my_mysql
from quant_researcher.quant.project_tool.logger.my_logger import LOG

T_STOCK_INDUSTRY_CLS = 'pub_inducls_components'
T_STOCK_INFO = 'astk_bd_stockinfo'
T_STOCK_INFO_LONGBO = 'jl_base_1102'


def get_stock_sname(stock_code=None, stock_type=None):
    """
    获取股票中文简称

    :param str or list or None stock_code: 股票代码，支持单个(str)多个(list)全部(None)
    :param str or list or None stock_type: 股票类型，可选：ES0101001-A股，ES0101002-B股，ES0102001001-H股，
                                                ES0102001-港股，ES0102002003-红筹股
    :return: pd.DataFrame
    """
    conn = db_conn.get_stock_conn()
    where = []
    if isinstance(stock_code, str):
        where.append(f"sec_code = '{stock_code}'")
    elif isinstance(stock_code, list):
        where.append(get_lst_condition(stock_code, "sec_code"))
    elif stock_code is None:
        pass
    else:
        LOG.error('目前stock_code只支持str、list或者None，请按要求输入')
        raise NotImplementedError

    if isinstance(stock_type, str):
        where.append(f"sec_type = '{stock_type}'")
    elif isinstance(stock_type, list):
        where.append(get_lst_condition(stock_type, "sec_type"))
    elif stock_type is None:
        where.append(get_lst_condition(['ES0101001', 'ES0101002', 'ES0102001001', 'ES0102001', 'ES0102002003'],
                                       "sec_type"))
    else:
        LOG.error('目前stock_type只支持str、list或者None，请按要求输入')
        raise NotImplementedError

    df = my_mysql.read_v2(select=['sec_code', 'sec_sname'], where=where, sfrom=T_STOCK_INFO,
                          conn=conn)
    conn.close()

    if df.empty:
        LOG.error(f"没有找到{stock_code}的中文简称，请检查输入")
        return

    df.rename(columns={'sec_code': "stock_code", 'sec_sname': "stock_sname"}, inplace=True)

    return df


def get_stock_base_info(stock_code=None, stock_type=None, select=None):
    """
    得到股票的基本信息

    :param str or list or None stock_code: 股票代码，支持单个(str)多个(list)全部(None)
    :param str or list or None stock_type: 股票类型，可选：ES0101001-A股，ES0101002-B股，ES0102001001-H股，
                                                ES0102001-港股，ES0102002003-红筹股
    :param list select: 需要返回哪些列，None则返回全部
    :return: pd.DataFrame
                - stock_code: 股票代码
                - sec_innercode: 股票内部代码
                - sec_name: 证券中文名称
                - sec_sname: 证券中文简称
                - sec_enname: 公司英文全称
                - sec_ensname: 公司英文简称
                - sec_type: 证券类型
                - hqmkt_code: 行情市场代码 沪市:SH,深市:SZ,沪港通:HK 深港通:SK,股转:SZ
                - str listing_state: 上市状态 1=正常上市,2=预上市,3=暂停,4=上市失败,5=终止,9=其他,10=交易,11=停牌,12=摘牌
                - listed_date: 上市时间 str "20200801"
                - delisted_date: 退市时间 str "20200801"
                - listed_board: 上市板块 1001=主板,1002=中小企业板,1003=创业板,1004=科创板,1005=三板,1006=大宗交易系统,1099=其他
                - int is_st: 是否ST股票 0=不是,1=是
                - comp_code: 归属公司编码
    """
    conn = db_conn.get_stock_conn()
    where = []
    if isinstance(stock_code, str):
        where.append(f"sec_code = '{stock_code}'")
    elif isinstance(stock_code, list):
        where.append(get_lst_condition(stock_code, "sec_code"))
    elif stock_code is None:
        pass
    else:
        LOG.error('目前stock_code只支持str、list或者None，请按要求输入')
        raise NotImplementedError

    if isinstance(stock_type, str):
        where.append(f"sec_type = '{stock_type}'")
    elif isinstance(stock_type, list):
        where.append(get_lst_condition(stock_type, "sec_type"))
    elif stock_type is None:
        where.append(get_lst_condition(['ES0101001', 'ES0101002', 'ES0102001001', 'ES0102001', 'ES0102002003'],
                                       "sec_type"))
    else:
        LOG.error('目前stock_type只支持str、list或者None，请按要求输入')
        raise NotImplementedError

    if select is not None:
        select = ['sec_code' if x == 'stock_code' else x for x in select]

    df = my_mysql.read_v2(select=select, where=where, sfrom=T_STOCK_INFO, conn=conn)
    conn.close()

    if df.empty:
        LOG.error(f"没有找到{stock_code}的基本信息，请检查输入")
        return

    if 'sec_code' in df.columns:
        df = df.rename(columns={'sec_code': 'stock_code'})

    if select is None:
        drop_column_list = ['tid', 'titime', 'tutime', 'tutimestamp',
                            'tcreator', 'tmodifier', 'tauditor', 'taudittime', 'tdatastate',
                            'tdatachannel', 'torigdatabuid']
        df.drop(columns=drop_column_list, inplace=True)

    return df


def get_alla_stock():
    """
    获取所有A股股票信息

    :return:
    """
    stock_code_df = get_stock_base_info(stock_type='ES0101001')
    mask1 = (stock_code_df['listing_state'] == '1')
    mask2 = (stock_code_df['listed_board'].isin(['1001', '1002', '1003', '1004']))
    stock_code_df = stock_code_df[mask1 & mask2]

    return stock_code_df


def get_terminated_stock():
    """
    获取终止上市股票列表
    Return
    --------
    DataFrame
        code :股票代码
        name :股票名称
        oDate:上市日期
        tDate:终止上市日期
    """
    try:

        ref = ct.SSEQ_CQ_REF_URL % (ct.P_TYPE['http'], ct.DOMAINS['sse'])
        clt = Client(rv.TERMINATED_URL % (ct.P_TYPE['http'], ct.DOMAINS['sseq'],
                                          ct.PAGES['ssecq'], _random(5),
                                          _random()), ref=ref, cookie=rv.MAR_SH_COOKIESTR)
        lines = clt.gvalue()
        lines = lines.decode('utf-8') if ct.PY3 else lines
        lines = lines[19:-1]
        lines = json.loads(lines)
        df = pd.DataFrame(lines['result'], columns=rv.TERMINATED_T_COLS)
        df.columns = rv.TERMINATED_COLS
        return df
    except Exception as er:
        print(str(er))


def get_suspended_stock():
    """
    获取暂停上市股票列表
    Return
    --------
    DataFrame
        code :股票代码
        name :股票名称
        oDate:上市日期
        tDate:终止上市日期
    """
    try:

        ref = ct.SSEQ_CQ_REF_URL % (ct.P_TYPE['http'], ct.DOMAINS['sse'])
        clt = Client(rv.SUSPENDED_URL % (ct.P_TYPE['http'], ct.DOMAINS['sseq'],
                                         ct.PAGES['ssecq'], _random(5),
                                         _random()), ref=ref, cookie=rv.MAR_SH_COOKIESTR)
        lines = clt.gvalue()
        lines = lines.decode('utf-8') if ct.PY3 else lines
        lines = lines[19:-1]
        lines = json.loads(lines)
        df = pd.DataFrame(lines['result'], columns=rv.TERMINATED_T_COLS)
        df.columns = rv.TERMINATED_COLS
        return df
    except Exception as er:
        print(str(er))


def get_st_stock():
    """
        获取风险警示板股票
    Return
    --------
    DataFrame
        code :股票代码
        name :股票名称
    """
    df = fd.get_stock_basics()
    df.reset_index(inplace=True)
    df = df[ct.FOR_CLASSIFY_COLS]
    df = df.loc[df.name.str.contains('ST')]
    df = df.sort_values('code').reset_index(drop=True)
    return df


def get_gem_stock():
    """
        获取创业板股票
    Return
    --------
    DataFrame
        code :股票代码
        name :股票名称
    """
    df = fd.get_stock_basics()
    df.reset_index(inplace=True)
    df = df[ct.FOR_CLASSIFY_COLS]
    df = df.loc[df.code.str[0] == '3']
    df = df.sort_values('code').reset_index(drop=True)
    return df


def get_sme_stock():
    """
        获取中小板股票
    Return
    --------
    DataFrame
        code :股票代码
        name :股票名称
    """
    df = fd.get_stock_basics()
    df.reset_index(inplace=True)
    df = df[ct.FOR_CLASSIFY_COLS]
    df = df.loc[df.code.str[0:3] == '002']
    df = df.sort_values('code').reset_index(drop=True)
    return df


def get_code_relations_longbo():
    """
    获得龙柏表中股票代码和公司代码之间的对应关系
    龙柏采用自己的一套编码，且有些数据我们表里还没有，需要读取龙柏表里的数据

    :return: pd.DataFrame
    """
    conn = db_conn.get_third_party_conn()
    select = ['distinct COMCODE as stockcode', 'STOCKCODE as code', 'CODE as code_sec']
    # F1102=1表示A股
    where = ["left(STOCKCODE, 1) in ('0','3','6')", "F1102_022 = '1'"]
    df = my_mysql.read_v2(select=select, sfrom=T_STOCK_INFO_LONGBO, where=where, conn=conn)
    conn.close()

    if df.empty:
        LOG.error(f"{T_STOCK_INFO_LONGBO}该表可能为空，请检查")
        return
    return df


def get_stock_industry_classification(induclsstd_code='sw2014', indu_code=None,
                                      indu_name=None, end_date=None, **kwargs):
    """
    获取股票行业分类

    :param str induclsstd_code: 支持'sw2014'申万行业分类2014版, 'csrc2012'证监会行业分类标准2012版
                                   'citic2019'中信行业2019分类
    :param list or str indu_code: 各行业分类下行业代码, 可参考表pub_inducls_define
    :param list or str indu_name: 各行业分类下行业名称，可参考表pub_inducls_define
    :param str or datetime.date or datetime.datetime or int end_date: 查询时间，
            支持带杠字符串、不带杠字符串、datetime、date、int格式，默认为None，即获取最新的行业分类数据
    :param kwargs:
        - conn: 数据库连接
    :return:
    """
    conn = kwargs.pop('conn', None)
    tmp_conn = db_conn.get_basic_data_conn() if conn is None else conn

    if end_date is not None:
        end_date = time_tool.format_date_str(end_date, '%Y%m%d')
        if end_date < '20140101' and induclsstd_code == 'sw2014':
            induclsstd_code = 'sw'
        end_date_condition = f"((classify_state = 1 and info_pubdate <= '{end_date}') " \
                             f"or (classify_state = 0 and info_pubdate < '{end_date}' " \
                             f"and out_date > '{end_date}')) "
    else:
        end_date_condition = f"classify_state = 1"

    where_condition = f"induclsstd_code = '{induclsstd_code}'"
    where_condition = where_condition + ' and ' + end_date_condition
    if indu_code is not None:
        if isinstance(indu_code, str):
            indu_code = [indu_code]
        indu_code_condition = get_lst_condition(indu_code, 'classify_code1')
        where_condition = where_condition + ' and ' + indu_code_condition
    if indu_name is not None:
        if isinstance(indu_name, str):
            indu_name = [indu_name]
        indu_name_condition = get_lst_condition(indu_name, 'classify_name1')
        where_condition = where_condition + ' and ' + indu_name_condition

    if induclsstd_code == 'sw':
        stock_industry = my_mysql.read_v2(select=['secu_code as stock_code', 'classify_code1 as indu_code',
                                                  'classify_code2', 'classify_code3', 'classify_name1 as indu_name'],
                                          sfrom=T_STOCK_INDUSTRY_CLS, where=where_condition, conn=tmp_conn)
        if conn is None:
            tmp_conn.close()
        if stock_industry.empty:
            LOG.error(f"没有查询到{induclsstd_code}的行业分类数据，请检查")
            return
        cls_2 = stock_industry['classify_code2'].where(stock_industry['classify_code2'].isin(SW_OLD_NEW_MAPPING.keys())).replace(SW_OLD_NEW_MAPPING)
        cls_3 = stock_industry['classify_code3'].where(stock_industry['classify_code3'].isin(SW_OLD_NEW_MAPPING.keys())).replace(SW_OLD_NEW_MAPPING)
        cls_2.loc[cls_3.notna()] = cls_3
        stock_industry['indu_code'] = cls_2.fillna(stock_industry['indu_code'])
        stock_industry['indu_name'] = stock_industry['indu_code'].replace(SW_INDUSTRY_CODE_SNAME)
        stock_industry = stock_industry[['stock_code', 'indu_code', 'indu_name']]
    else:
        stock_industry = my_mysql.read_v2(select=['secu_code as stock_code', 'classify_code1 as indu_code',
                                                  'classify_name1 as indu_name'],
                                          sfrom=T_STOCK_INDUSTRY_CLS, where=where_condition, conn=tmp_conn)
        if conn is None:
            tmp_conn.close()
        if stock_industry.empty:
            LOG.error(f"没有查询到{induclsstd_code}的行业分类数据，请检查")
            return
    # 剔除非A股股票
    stock_industry = stock_industry[stock_industry.stock_code.str[0].isin(['0', '3', '6'])]

    return stock_industry


if __name__ == '__main__':
    # 测试
    aaa = get_stock_sname()
    # aaa = get_stock_base_info(['000001', '000002'], select=['stock_code', 'is_st'])
    test = get_stock_industry_classification(induclsstd_code='sw2014', end_date='2013-12-01')